# -----------------------------------------------------------------------------
# Copyright (c) 2025, Southeast University (China)
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#  - Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
#  - Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
#  - Neither the name of the copyright holders nor the names of its
#    contributors may be used to endorse or promote products derived from
#    this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# Authors: Jiajie Xu, Yanfang Zhang, Jiaqi Gao, Leyun Tian
# -----------------------------------------------------------------------------


from Algo_var import *
import networkx as nx
from TimingGraph import TimingGraph
from TimingLib import TimingLib
# import concurrent
# from concurrent.futures import ProcessPoolExecutor
import work_var
from Delay_Calculator import Delay_Calculator
import numpy as np
from tqdm import tqdm
import copy
import Lagrange_Solver


# ALGO LOGIC: initialize Timing propagation algorithm here:
class Timing_Propagation(TimingGraph, TimingLib, Delay_Calculator):
    """Timing Propogation method for test design"""

    """
    self.design -> current design name
    self.corner -> tuple of current process corner
    self.rc_corner -> string of current RC corner
    self.G -> processed timing graph with timing arcs of only positive_unate, negative_unate and rising_edge
    self.Lagrange_Solver -> a Lagrange solver to calculate lamba and select cell
    Newly extended Graph features:
        nodes:
            'arrival_pt': golden arrival time from PT
            'trans_pt': golden transition time from PT
            'slack_pt': golden slack from PT
            'ceff_pt': golden ceff from PT
            'ECO_visited': a flag in ECO mode to adjust if the node is visited before
            'lambda': a lagrange factor to measure the importance of the edge
        edges:
            'delay_pt': golden delay from PT
            'trans': transition time of the timing arc
    function:
    self.Test_Update_timing(self) -> test propagete timing from PI to PO using golden PT data
    self.ECO_Iteration(self, max_iteration, arnoldi, Simple) -> Conduct ECO iteration for the design
    self.Forward_Propagation_Simple(self, ECO) -> Simple forward propagate timing from PI to PO,
                                                total capacitance is utilized for cell delay calculation,
                                                if ECO, generate ECO solution for netlist,
    self._Calc_Cell_Timing_Simple(self, node_arrival, node_trans, node, predecessor, key, edge_data, ECO, local_cells) ->
                                                    Simply calculate cell arc NLDM Timing in self.G or ECO subGraph
                                                    local_cells is a dict represent local subgraph cell refname to cellname
    self.Forward_Propagation(self, ECO, arnoldi) -> forward propagate timing from PI to PO,
                                                    if ECO, generate ECO solution for netlist,
                                                    if Arnoldi is False, using Elmore net delay
    self._Calc_Cell_Timing(self, node_arrival, node_trans, node, predecessor, key, edge_data, arnoldi, ECO, local_cells) ->
                                                    Calculate cell arc NLDM Timing in self.G or ECO subGraph
                                                    local_cells is a dict represent local subgraph cell refname to cellname
    self._Update_parasitic_(self, parasitic, rf) -> Update parasitic based on NLDM r/f receiver caps
    self.Estimate_ECO_Simple(self, target_node, ori_cell, new_cell) -> Simply Estimate Subgraph Timing after ECO
    self.Estimate_ECO(self, target_node, ori_cell, new_cell, arnoldi = True) -> Estimate Subgraph Timing after ECO
    self._Extract_Local_Subgraph(self, node) -> make a copy of ECO subgraph for node
    self.Implement_ECO(self, Copied_G) -> Implement ECO  subgraph to self.G
    self.Backward_Propagation(self) -> Propagate slack from PO to PI
    self.Report_Global_Timing(self) -> A simple report for total negative slacks(TNS) and worst negative slack(WNS) on EPs
    """

    def __init__(
        self,
        design: str = "mc_top",
        contrast: bool = False,
        clock_period: float = 0.5,
        corner: str = "ssgnp",  # process corner from: ssgnp, ffgnp and tt
        temperature: str = "125",  # temperature from: m40, 0, 125; ffgnp: m40, 0, 125; tt: 25, 85
        voltage: str = "0p63",  # voltage from: ssgnp: ssgnp: 0p63, 0p72, 0p9; ffgnp: 0p77, 0p88, 1p05; tt: 0p7, 0p8, 1
        rc_corner: str = "rcworst_CCworst",
    ):
        TimingGraph.__init__(self, design)
        TimingLib.__init__(self)
        Delay_Calculator.__init__(
            self, design, rc_corner, th_slew1=0.3, th_slew2=0.7, th_delay=0.5
        )
        self.design = design
        self.corner = (corner, voltage, temperature)
        self.rc_corner = rc_corner

        # mc_top 0.5, s27 0.1
        self.clk_period = clock_period
        """
        unnecessary timing arcs represents cell timing arcs which have timing sense unate other than negative_unate, positive_unate or rising_edge,
        or arcs which does not have delays
        These arcs belong to registers which are not considered for optimization
        """
        unnecessary_arcs = []
        for u, v, k, data in self.G.edges(keys=True, data=True):
            if data["is_cell"]:
                if (
                    data["sense_unate"] != "positive_unate"
                    and data["sense_unate"] != "negative_unate"
                    and data["sense_unate"] != "rising_edge"
                    and data["sense_unate"] != "falling_edge"
                    and data["sense_unate"] != "setup_clk_rise"
                    and data["sense_unate"] != "hold_clk_rise"
                    and data["sense_unate"] != "setup_clk_fall"
                    and data["sense_unate"] != "hold_clk_fall"
                ):
                    unnecessary_arcs.append((u, v, k))
            if all(x is None for x in data["delay"]):
                unnecessary_arcs.append((u, v, k))
        self.G.remove_edges_from(unnecessary_arcs)

        try:
            # topologically sort nodes, isolated nodes are not included
            self.levelized_nodes = list(nx.topological_sort(self.G))
            self.reversed_levelized_nodes = list(nx.topological_sort(self.G))[::-1]
        except Exception as e:
            print(f"Error occurred: {e}\n Timing Graph cannot be levelized.")

        self.Lagrange_Solver = Lagrange_Solver.Lagrange_Solver(contrast=contrast)

        # save the timing result from pt for validation and net delay/trans calculation
        print("Saving pin timing features for validation and initialize ECO flag")
        for node in self.levelized_nodes:
            self.G.nodes[node]["arrival_pt"] = self.G.nodes[node]["arrival"].copy()
            self.G.nodes[node]["trans_pt"] = self.G.nodes[node]["trans"].copy()
            self.G.nodes[node]["slack_pt"] = self.G.nodes[node]["slack"].copy()
            self.G.nodes[node]["ceff_pt"] = self.G.nodes[node]["ceff"].copy()
            self.G.nodes[node]["ECO_visited"] = False

        print("Saving arc timing features for validation")
        for u, v, k in self.G.edges(keys=True, data=False):
            self.G[u][v][k]["delay_pt"] = self.G[u][v][k]["delay"].copy()
            self.G[u][v][k]["trans"] = [None, None, None, None]

    def ECO_Iteration(self, max_iteration=5, arnoldi=True, Simple=False):
        if not Simple:
            self.Forward_Propagation(ECO=False, arnoldi=arnoldi)
        else:
            self.Forward_Propagation_Simple(ECO=False)
        self.Backward_Propagation()
        print(f"Initialize timing info.")
        self.Report_Global_Timing()

        print(f"Begin ECO optimization.")
        for iter in range(max_iteration):
            if not Simple:
                self.Forward_Propagation(ECO=True, arnoldi=arnoldi)
            else:
                self.Forward_Propagation_Simple(ECO=True)
            self.Backward_Propagation()
            print(f"{iter}th iteration finished.")
            self.Report_Global_Timing()
            if not self.Lagrange_Solver.Changes_Made:
                print("No changes made this iteration, break earlier.")
                break
        print(f"Finished ECO optimization iteration with maximum {iter + 1} iteration.")
        print(f"Output ECO changelist for PrimeTime.")
        self.Lagrange_Solver.Output_ECO_Change_List()

    def Test_Forward_Propagation(self):
        print("Using PT golden ceff and net delay to calculate forward propagation.")
        for node in self.levelized_nodes:
            # if item in list are all None, ignore the pin
            if all(x is None for x in self.G.nodes[node]["arrival"]):
                continue
            in_edges = list(self.G.in_edges(node, data=True, keys=True))
            # Primary Inputs are not updated
            if not in_edges:
                pass
            # initialize timing features
            else:
                node_arrival = [
                    -float("inf"),
                    -float("inf"),
                    float("inf"),
                    float("inf"),
                ]
                node_trans = [-float("inf"), -float("inf"), float("inf"), float("inf")]
                # traverse input nodes
                for predecessor, _, k, edge_data in self.G.in_edges(
                    node, data=True, keys=True
                ):
                    pred_arrival = self.G.nodes[predecessor][
                        "arrival"
                    ]  # Max_Rise, Max_Fall, Min_Rise, Min_Fall
                    pred_trans = self.G.nodes[predecessor][
                        "trans"
                    ]  # Max_Rise, Max_Fall, Min_Rise, Min_Fall
                    if edge_data.get("is_cell"):
                        # if the timing arc is a cell arc, using NLDM model for delay and trans calculation
                        # (arc,rf,sense,type,when,sdf_cond)
                        arc = (
                            predecessor.split("/")[1] + "->" + node.split("/")[1]
                        )  # get timing arc
                        cellname = self.cells[
                            edge_data["name"]
                        ]  # get cell name in based on the refname
                        node_ceff = self.G.nodes[node][
                            "ceff"
                        ]  # Max_Rise, Max_Fall, Min_Rise, Min_Fall
                        unate = edge_data["sense_unate"]
                        when = edge_data["when"]
                        sdf_cond = edge_data["sdf_cond"]
                        # nomal logical arc
                        if unate != "rising_edge":
                            type = "combinational"
                            for i, out_rf in enumerate(["r", "f"]):
                                if unate == "positive_unate":
                                    j = i
                                else:
                                    j = 1 - i
                                # in_rf is relay on unate and out_rf
                                delay_max, out_trans_max = self.NLDM_Timing(
                                    self.corner,
                                    cellname,
                                    arc,
                                    out_rf,
                                    unate,
                                    type,
                                    when,
                                    sdf_cond,
                                    pred_trans[j],
                                    node_ceff[i],
                                )
                                delay_min, out_trans_min = self.NLDM_Timing(
                                    self.corner,
                                    cellname,
                                    arc,
                                    out_rf,
                                    unate,
                                    type,
                                    when,
                                    sdf_cond,
                                    pred_trans[j + 2],
                                    node_ceff[i + 2],
                                )

                                arrival_max = delay_max + pred_arrival[j]
                                arrival_min = delay_min + pred_arrival[j + 2]

                                self.G[predecessor][node][k]["delay"][i] = delay_max
                                self.G[predecessor][node][k]["delay"][i + 2] = delay_min

                                node_arrival[i] = max(node_arrival[i], arrival_max)
                                node_arrival[i + 2] = min(
                                    node_arrival[i + 2], arrival_min
                                )
                                node_trans[i] = max(node_trans[i], out_trans_max)
                                node_trans[i + 2] = min(
                                    node_trans[i + 2], out_trans_min
                                )
                        else:  # rising_edge unate, using non_unate in liberty
                            type = unate
                            unate = "non_unate"

                            delay_max_rise, out_trans_max_rise = self.NLDM_Timing(
                                self.corner,
                                cellname,
                                arc,
                                "r",
                                unate,
                                type,
                                when,
                                sdf_cond,
                                pred_trans[0],
                                node_ceff[0],
                            )  # input trans Max_Rise, output ceff Max_Rise
                            delay_max_fall, out_trans_max_fall = self.NLDM_Timing(
                                self.corner,
                                cellname,
                                arc,
                                "f",
                                unate,
                                type,
                                when,
                                sdf_cond,
                                pred_trans[0],
                                node_ceff[1],
                            )  # input trans Max_Rise, output ceff Max_Fall
                            delay_min_rise, out_trans_min_rise = self.NLDM_Timing(
                                self.corner,
                                cellname,
                                arc,
                                "r",
                                unate,
                                type,
                                when,
                                sdf_cond,
                                pred_trans[2],
                                node_ceff[2],
                            )  # input trans Min_Rise, output ceff Min_Rise
                            delay_min_fall, out_trans_min_fall = self.NLDM_Timing(
                                self.corner,
                                cellname,
                                arc,
                                "f",
                                unate,
                                type,
                                when,
                                sdf_cond,
                                pred_trans[2],
                                node_ceff[3],
                            )  # input trans Min_Rise, output ceff Min_Fall
                            arrival_max_rise = delay_max_rise + pred_arrival[0]
                            arrival_max_fall = delay_max_fall + pred_arrival[0]
                            arrival_min_rise = delay_min_rise + pred_arrival[2]
                            arrival_min_fall = delay_min_fall + pred_arrival[2]

                            self.G[predecessor][node][k]["delay"] = [
                                delay_max_rise,
                                delay_max_fall,
                                delay_min_rise,
                                delay_min_fall,
                            ]

                            node_arrival = [
                                max(node_arrival[0], arrival_max_rise),
                                max(node_arrival[1], arrival_max_fall),
                                min(node_arrival[2], arrival_min_rise),
                                min(node_arrival[3], arrival_min_fall),
                            ]
                            node_trans = [
                                max(node_trans[0], out_trans_max_rise),
                                max(node_trans[1], out_trans_max_fall),
                                min(node_trans[2], out_trans_min_rise),
                                min(node_trans[3], out_trans_min_fall),
                            ]
                    else:
                        # if the timing arc is not a cell arc, its a net arc
                        # use the timing data from PrimeaTime to calculate arrival time and transition
                        # adjust if None in predecessor to avoid clk port
                        edge_arrival = [
                            (a if a is not None else 0) + (b if b is not None else 0)
                            for a, b in zip(
                                self.G.nodes[predecessor]["arrival"],
                                self.G[predecessor][node][k]["delay"],
                            )
                        ]
                        edge_trans = [
                            a
                            - (b if b is not None else 0)
                            + (c if c is not None else 0)
                            for a, b, c in zip(
                                self.G.nodes[node]["trans_pt"],
                                self.G.nodes[predecessor]["trans_pt"],
                                self.G.nodes[predecessor]["trans"],
                            )
                        ]
                        # update node timing
                        for idx in [0, 1]:  # Max_Rise, Max_Fall
                            node_arrival[idx] = max(
                                node_arrival[idx], edge_arrival[idx]
                            )
                            node_trans[idx] = max(node_trans[idx], edge_trans[idx])
                        for idx in [2, 3]:  # Min_Rise, Min_Fall
                            node_arrival[idx] = min(
                                node_arrival[idx], edge_arrival[idx]
                            )
                            node_trans[idx] = min(node_trans[idx], edge_trans[idx])
                self.G.nodes[node]["arrival"] = node_arrival
                self.G.nodes[node]["trans"] = node_trans

    def Forward_Propagation_Simple(self, ECO=False):
        """Using total capacitance to calculate NLDM cell delay, net delay is calculated by Elmore model"""
        if ECO:
            # Update the lambda multipliers
            self.Lagrange_Solver.update_lagrange_multipliers(
                self.G, self.reversed_levelized_nodes
            )
            # Initialize ECO visited flag to False
            for node in self.levelized_nodes:
                self.G.nodes[node]["ECO_visited"] = False
        
        desc_text = "Forward ECO Propagation, Simple Mode" if ECO else "Forward Timing Propagation, Simple Mode"
        for node in tqdm(
            self.levelized_nodes,
            desc=desc_text,
            ncols=110,
            colour="cyan",
            bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [Elapsed: {elapsed} | Speed: {rate_fmt}]",
            position=0,
        ):
        #for node in self.levelized_nodes:
            # if item in list are all None, ignore the pin
            if all(x is None for x in self.G.nodes[node]["arrival"]):
                continue
            in_edges = list(self.G.in_edges(node, data=True, keys=True))
            # Primary Inputs are not updated
            if not in_edges:
                pass
            # initialize timing features
            else:
                node_arrival = [
                    -float("inf"),
                    -float("inf"),
                    float("inf"),
                    float("inf"),
                ]
                node_trans = [
                    -float("inf"),
                    -float("inf"),
                    float("inf"), 
                    float("inf")
                    ]
                # traverse input nodes
                for predecessor, _, key, edge_data in self.G.in_edges(
                    node, data=True, keys=True
                ):
                    if ECO:
                        # the node is visited
                        self.G.nodes[predecessor]["ECO_visited"] = True
                    if (
                        edge_data.get("is_cell")
                        and not edge_data["sense_unate"].startswith("setup")
                        and not edge_data["sense_unate"].startswith("hold")
                    ):
                        self._Calc_Cell_Timing_Simple(
                            node_arrival,
                            node_trans,
                            node,
                            predecessor,
                            key,
                            edge_data,
                            ECO=False,
                            local_cells=None,
                        )    
                    elif(
                        not edge_data.get("is_cell")
                    ):
                        # # if the timing arc is not a cell arc, its a net arc
                        # # Elmore to calculate cell delay
                        # pred_arrival = self.G.nodes[predecessor][
                        #     "arrival"
                        # ]  # Max_Rise, Max_Fall, Min_Rise, Min_Fall
                        # pred_trans = self.G.nodes[predecessor][
                        #     "trans"
                        # ]  # Max_Rise, Max_Fall, Min_Rise, Min_Fall
                        # # not update timing with None in predecessor
                        # if None in pred_arrival:
                        #     node_arrival = self.G.nodes[node]["arrival"]
                        #     node_trans = self.G.nodes[node]["trans"]
                        # else:
                        #     for i, (ap, tp) in enumerate(
                        #         zip(pred_arrival, pred_trans)
                        #     ):
                        #         delay, trans = self.Calc_Elmore(
                        #             edge_data.get("name"), tp, node
                        #         )
                        #         self.G[predecessor][node][key]["delay"][i] = delay
                        #         self.G[predecessor][node][key]["trans"][i] = trans
                        #         if i < 2:
                        #             node_arrival[i] = max(
                        #                 node_arrival[i], ap + delay
                        #             )
                        #             node_trans[i] = max(node_trans[i], trans)
                        #         else:
                        #             node_arrival[i] = min(
                        #                 node_arrival[i], ap + delay
                        #             )
                        #             node_trans[i] = min(node_trans[i], trans)
                        
                        # if the timing arc is not a cell arc, its a net arc
                        # use the timing data from PrimeaTime to calculate arrival time and transition
                        # adjust if None in predecessor to avoid clk port
                        edge_arrival = [
                            (a if a is not None else 0) + (b if b is not None else 0)
                            for a, b in zip(
                                self.G.nodes[predecessor]["arrival"],
                                self.G[predecessor][node][key]["delay"],
                            )
                        ]
                        edge_trans = [
                            a
                            - (b if b is not None else 0)
                            + (c if c is not None else 0)
                            for a, b, c in zip(
                                self.G.nodes[node]["trans_pt"],
                                self.G.nodes[predecessor]["trans_pt"],
                                self.G.nodes[predecessor]["trans"],
                            )
                        ]
                        # update node timing
                        for idx in [0, 1]:  # Max_Rise, Max_Fall
                            node_arrival[idx] = max(
                                node_arrival[idx], edge_arrival[idx]
                            )
                            node_trans[idx] = max(node_trans[idx], edge_trans[idx])
                        for idx in [2, 3]:  # Min_Rise, Min_Fall
                            node_arrival[idx] = min(
                                node_arrival[idx], edge_arrival[idx]
                            )
                            node_trans[idx] = min(node_trans[idx], edge_trans[idx])
                        
                self.G.nodes[node]["arrival"] = node_arrival
                self.G.nodes[node]["trans"] = node_trans
                
                if (
                    ECO
                    and edge_data.get("is_cell")
                    and not self.G.nodes[predecessor]["is_clock_network"]
                ):
                    cellname = self.cells[edge_data["name"]]
                    Estimated_subgraph = {}
                    # Estimated current sub graph timing
                    Estimated_subgraph[cellname] = self.Estimate_ECO_Simple(
                        node, cellname, None
                    )
                    # Estimated all candidate sub graph timing
                    candidate_cells = self.Get_Candidate(cellname)
                    for cell in candidate_cells:
                        if cell != cellname:
                            Estimated_subgraph[cell] = self.Estimate_ECO_Simple(
                                node, cellname, cell
                            )
                    
                    # call Lagrange Solver to generate solution
                    selected_cell = self.Lagrange_Solver.Select_Cell(
                        Estimated_subgraph,
                        cellname,
                        candidate_cells,
                        node.split("/")[0],
                    )

                    # Update Timing Graph and self.cells with selected_cell
                    self.Implement_ECO(Estimated_subgraph[selected_cell])
                    self.cells[edge_data["name"]] = selected_cell

                    # Update all affected parts in timing graph
                    # traverse all cell inputs
                    for pred1, _, key1, edge_data1 in self.G.in_edges(
                        node, data=True, keys=True
                    ):
                        # all predecessor cell outputs
                        for pred2, _, key2, edge_data2 in self.G.in_edges(
                            pred1, data=True, keys=True
                        ):
                            visited = []
                            # successor cell inputs
                            for _, succ1, _, _ in self.G.out_edges(
                                pred2, data=True, keys=True
                            ):
                                visited.append(self.G.nodes[succ1]["ECO_visited"])
                            # if all succ1 are already visited, update them together
                            if all(visited):
                                for _, succ1, key3, edge_data3 in self.G.out_edges(
                                    pred2, data=True, keys=True
                                ):
                                    # not update the already newest ECO cell timing
                                    if succ1 != pred1:
                                        for (
                                            _,
                                            succ2,
                                            key4,
                                            edge_data4,
                                        ) in self.G.out_edges(
                                            succ1, data=True, keys=True
                                        ):
                                            succ2_arrival = self.G.nodes[succ2][
                                                "arrival"
                                            ]
                                            succ2_trans = self.G.nodes[succ2][
                                                "trans"
                                            ]
                                            self._Calc_Cell_Timing_Simple(
                                                succ2_arrival,
                                                succ2_trans,
                                                succ2,
                                                succ1,
                                                key4,
                                                edge_data4,
                                                ECO=False,
                                                local_cells=None,
                                            )
                                            self.G.nodes[succ2][
                                                "arrival"
                                            ] = succ2_arrival
                                            self.G.nodes[succ2][
                                                "trans"
                                            ] = succ2_trans
    
    def _Calc_Cell_Timing_Simple(
        self,
        node_arrival,
        node_trans,
        node,
        predecessor,
        key,
        edge_data,
        ECO=False,
        local_graph=None,
        local_cells=None,
    ):
        """
        This is the simple version of Calculate cell timing, Usining NLDM and total capacitance
        node_arrival: the arrival time to be calculated
        node_trans: the transition time to be calculated
        node: the node name
        predecessor: the predecessor name
        key: edge key
        edge_data: Graph edge data
        ECO: if in ECO mode
        local_cells: the local cell dict for the copied graph
        """
        if ECO:
            assert local_graph is not None, "Local graph is not given"
            assert local_cells is not None, "Local cell dict is not given"
            Graph = local_graph
        else:
            Graph = self.G
        pred_arrival = Graph.nodes[predecessor][
            "arrival"
        ]  # Max_Rise, Max_Fall, Min_Rise, Min_Fall
        pred_trans = Graph.nodes[predecessor][
            "trans"
        ]  # Max_Rise, Max_Fall, Min_Rise, Min_Fall
        # (arc,rf,sense,type,when,sdf_cond)
        arc = (
            predecessor.split("/")[1] + "->" + node.split("/")[1]
        )  # get timing arc
        if ECO:
            cellname = local_cells[
                edge_data["name"]
            ]
        else:
            cellname = self.cells[
                edge_data["name"]
            ]  # get cell name in based on the refname
            
        fanout = [
            (v, k, edge_data)
            for _, v, k, edge_data in Graph.out_edges(node, data=True, keys=True)
        ]
        fanout_net = fanout[0][2]["name"]
        
        # prevent changes to self.Parasitics
        parasitic = copy.deepcopy(self.Parasitics[fanout_net])
        parasitic = self._Update_parasitic_(parasitic, "r", ECO, local_cells)
        total_cap_r = parasitic._total_cap()
        
        parasitic = copy.deepcopy(self.Parasitics[fanout_net])
        parasitic = self._Update_parasitic_(parasitic, "f", ECO, local_cells)
        total_cap_f = parasitic._total_cap()
        
        node_ceff = [total_cap_r, total_cap_f, total_cap_r, total_cap_f]  # Max_Rise, Max_Fall, Min_Rise, Min_Fall

        unate = edge_data["sense_unate"]
        when = edge_data["when"]
        sdf_cond = edge_data["sdf_cond"]
        # nomal logical arc
        if unate != "rising_edge":
            type = "combinational"
            for i, out_rf in enumerate(["r", "f"]):
                if unate == "positive_unate":
                    j = i
                else:
                    j = 1 - i
                
                # in_rf is relay on unate and out_rf
                delay_max, out_trans_max = self.NLDM_Timing(
                    self.corner,
                    cellname,
                    arc,
                    out_rf,
                    unate,
                    type,
                    when,
                    sdf_cond,
                    pred_trans[j],
                    node_ceff[i],
                )
                delay_min, out_trans_min = self.NLDM_Timing(
                    self.corner,
                    cellname,
                    arc,
                    out_rf,
                    unate,
                    type,
                    when,
                    sdf_cond,
                    pred_trans[j + 2],
                    node_ceff[i + 2],
                )

                arrival_max = delay_max + pred_arrival[j]
                arrival_min = delay_min + pred_arrival[j + 2]

                Graph[predecessor][node][key]["delay"][i] = delay_max
                Graph[predecessor][node][key]["delay"][i + 2] = delay_min

                node_arrival[i] = max(
                    node_arrival[i], 
                    arrival_max
                )
                node_arrival[i + 2] = min(
                    node_arrival[i + 2], arrival_min
                )
                node_trans[i] = max(
                    node_trans[i], 
                    out_trans_max
                )
                node_trans[i + 2] = min(
                    node_trans[i + 2], out_trans_min
                )
        else:  # rising_edge unate, using non_unate in liberty
            type = unate
            unate = "non_unate"

            delay_max_rise, out_trans_max_rise = self.NLDM_Timing(
                self.corner,
                cellname,
                arc,
                "r",
                unate,
                type,
                when,
                sdf_cond,
                pred_trans[0],
                node_ceff[0],
            )  # input trans Max_Rise, output ceff Max_Rise
            delay_max_fall, out_trans_max_fall = self.NLDM_Timing(
                self.corner,
                cellname,
                arc,
                "f",
                unate,
                type,
                when,
                sdf_cond,
                pred_trans[0],
                node_ceff[1],
            )  # input trans Max_Rise, output ceff Max_Fall
            delay_min_rise, out_trans_min_rise = self.NLDM_Timing(
                self.corner,
                cellname,
                arc,
                "r",
                unate,
                type,
                when,
                sdf_cond,
                pred_trans[2],
                node_ceff[2],
            )  # input trans Min_Rise, output ceff Min_Rise
            delay_min_fall, out_trans_min_fall = self.NLDM_Timing(
                self.corner,
                cellname,
                arc,
                "f",
                unate,
                type,
                when,
                sdf_cond,
                pred_trans[2],
                node_ceff[3],
            )  # input trans Min_Rise, output ceff Min_Fall
            arrival_max_rise = delay_max_rise + pred_arrival[0]
            arrival_max_fall = delay_max_fall + pred_arrival[0]
            arrival_min_rise = delay_min_rise + pred_arrival[2]
            arrival_min_fall = delay_min_fall + pred_arrival[2]

            Graph[predecessor][node][key]["delay"] = [
                delay_max_rise,
                delay_max_fall,
                delay_min_rise,
                delay_min_fall,
            ]

            node_arrival[0] = max(node_arrival[0], arrival_max_rise)
            node_arrival[1] = max(node_arrival[1], arrival_max_fall)
            node_arrival[2] = min(node_arrival[2], arrival_min_rise)
            node_arrival[3] = min(node_arrival[3], arrival_min_fall)

            node_trans[0] = max(node_trans[0], out_trans_max_rise)
            node_trans[1] = max(node_trans[1], out_trans_max_fall)
            node_trans[2] = min(node_trans[2], out_trans_min_rise)
            node_trans[3] = min(node_trans[3], out_trans_min_fall)
    
    def Forward_Propagation(self, ECO=False, arnoldi=True):
        if ECO:
            # Update the lambda multipliers
            self.Lagrange_Solver.update_lagrange_multipliers(
                self.G, self.reversed_levelized_nodes
            )

            # Initialize ECO visited flag to False
            for node in self.levelized_nodes:
                self.G.nodes[node]["ECO_visited"] = False

        desc_text = "Forward Timing Propagation" if ECO else "Forward ECO Propagation"
        for node in tqdm(
            self.levelized_nodes,
            desc="Forward Timing Propagation",
            ncols=100,
            colour="cyan",
            bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [Elapsed: {elapsed} | Speed: {rate_fmt}]",
            position=0,
        ):
        # for node in self.levelized_nodes:
            # if item in list are all None, ignore the pin
            if all(x is None for x in self.G.nodes[node]["arrival"]):
                continue
            in_edges = list(self.G.in_edges(node, data=True, keys=True))
            # Primary Inputs are not updated
            if not in_edges:
                pass
            # initialize timing features
            else:
                node_arrival = [
                    -float("inf"),
                    -float("inf"),
                    float("inf"),
                    float("inf"),
                ]
                node_trans = [
                    -float("inf"),
                    -float("inf"),
                    float("inf"),
                    float("inf")
                    ]
                # traverse input nodes
                for predecessor, _, key, edge_data in self.G.in_edges(
                    node, data=True, keys=True
                ):
                    if ECO:
                        # the node is visited
                        self.G.nodes[predecessor]["ECO_visited"] = True
                    """
                    cell delay calculation, if alnordi, calculate success net delay after cell delay
                    """
                    if (
                        edge_data.get("is_cell")
                        and not edge_data["sense_unate"].startswith("setup")
                        and not edge_data["sense_unate"].startswith("hold")
                    ):
                        self._Calc_Cell_Timing(
                            node_arrival,
                            node_trans,
                            node,
                            predecessor,
                            key,
                            edge_data,
                            arnoldi,
                            ECO=False,
                            local_cells=None,
                        )
                    elif(
                        not edge_data.get("is_cell")
                    ):
                        """
                        Using Elmore model for net delay and trans calculation
                        """
                        if not arnoldi:
                            pred_arrival = self.G.nodes[predecessor][
                                "arrival"
                            ]  # Max_Rise, Max_Fall, Min_Rise, Min_Fall
                            pred_trans = self.G.nodes[predecessor][
                                "trans"
                            ]  # Max_Rise, Max_Fall, Min_Rise, Min_Fall
                            # not update timing with None in predecessor
                            if None in pred_arrival:
                                node_arrival = self.G.nodes[node]["arrival"]
                                node_trans = self.G.nodes[node]["trans"]
                            else:
                                for i, (ap, tp) in enumerate(
                                    zip(pred_arrival, pred_trans)
                                ):
                                    delay, trans = self.Calc_Elmore(
                                        edge_data.get("name"), tp, node
                                    )
                                    self.G[predecessor][node][key]["delay"][i] = delay
                                    self.G[predecessor][node][key]["trans"][i] = trans
                                    if i < 2:
                                        node_arrival[i] = max(
                                            node_arrival[i], ap + delay
                                        )
                                        node_trans[i] = max(node_trans[i], trans)
                                    else:
                                        node_arrival[i] = min(
                                            node_arrival[i], ap + delay
                                        )
                                        node_trans[i] = min(node_trans[i], trans)

                if edge_data.get("is_cell") or not arnoldi:
                    self.G.nodes[node]["arrival"] = node_arrival
                    self.G.nodes[node]["trans"] = node_trans

                if (
                    ECO
                    and edge_data.get("is_cell")
                    and not self.G.nodes[predecessor]["is_clock_network"]
                ):
                    cellname = self.cells[edge_data["name"]]
                    Estimated_subgraph = {}
                    # Estimated current sub graph timing
                    Estimated_subgraph[cellname] = self.Estimate_ECO(
                        node, cellname, None, arnoldi
                    )
                    # Estimated all candidate sub graph timing
                    candidate_cells = self.Get_Candidate(cellname)
                    for cell in candidate_cells:
                        if cell != cellname:
                            Estimated_subgraph[cell] = self.Estimate_ECO(
                                node, cellname, cell, arnoldi
                            )
                            
                    # call Lagrange Solver to generate solution
                    selected_cell = self.Lagrange_Solver.Select_Cell(
                        Estimated_subgraph,
                        cellname,
                        candidate_cells,
                        node.split("/")[0],
                    )

                    # Update Timing Graph and self.cells with selected_cell
                    self.Implement_ECO(Estimated_subgraph[selected_cell])
                    self.cells[edge_data["name"]] = selected_cell

                    # Update all affected parts in timing graph
                    # traverse all cell inputs
                    for pred1, _, key1, edge_data1 in self.G.in_edges(
                        node, data=True, keys=True
                    ):
                        # all predecessor cell outputs
                        for pred2, _, key2, edge_data2 in self.G.in_edges(
                            pred1, data=True, keys=True
                        ):
                            visited = []
                            # successor cell inputs
                            for _, succ1, _, _ in self.G.out_edges(
                                pred2, data=True, keys=True
                            ):
                                visited.append(self.G.nodes[succ1]["ECO_visited"])
                            # if all succ1 are already visited, update them together
                            if all(visited):
                                for _, succ1, key3, edge_data3 in self.G.out_edges(
                                    pred2, data=True, keys=True
                                ):
                                    # not update the already newest ECO cell timing
                                    if succ1 != pred1:
                                        for (
                                            _,
                                            succ2,
                                            key4,
                                            edge_data4,
                                        ) in self.G.out_edges(
                                            succ1, data=True, keys=True
                                        ):
                                            succ2_arrival = self.G.nodes[succ2][
                                                "arrival"
                                            ]
                                            succ2_trans = self.G.nodes[succ2]["trans"]
                                            self._Calc_Cell_Timing(
                                                succ2_arrival,
                                                succ2_trans,
                                                succ2,
                                                succ1,
                                                key4,
                                                edge_data4,
                                                arnoldi,
                                                ECO=False,
                                                local_cells=None,
                                            )
                                            self.G.nodes[succ2][
                                                "arrival"
                                            ] = succ2_arrival
                                            self.G.nodes[succ2][
                                                "trans"
                                            ] = succ2_trans

    def _Calc_Cell_Timing(
        self,
        node_arrival,
        node_trans,
        node,
        predecessor,
        key,
        edge_data,
        arnoldi=True,
        ECO=False,
        local_graph=None,
        local_cells=None,
    ):
        """
        node_arrival: the arrival time to be calculated
        node_trans: the transition time to be calculated
        node: the node name
        predecessor: the predecessor name
        key: edge key
        edge_data: Graph edge data
        arnoldi: if arnoldi mode utilized
        ECO: if in ECO mode
        local_cells: the local cell dict for the copied graph
        """
        if ECO:
            assert local_graph is not None, "Local graph is not given"
            assert local_cells is not None, "Local cell dict is not given"
            Graph = local_graph
        else:
            Graph = self.G
        pred_arrival = Graph.nodes[predecessor][
            "arrival"
        ]  # Max_Rise, Max_Fall, Min_Rise, Min_Fall
        pred_trans = Graph.nodes[predecessor][
            "trans"
        ]  # Max_Rise, Max_Fall, Min_Rise, Min_Fall
        arc = predecessor.split("/")[1] + "->" + node.split("/")[1]  # get timing arc
        if ECO:
            cellname = local_cells[
                edge_data["name"]
            ]
        else:
            cellname = self.cells[
                edge_data["name"]
            ]  # get cell name in based on the refname
        unate = edge_data["sense_unate"]
        when = edge_data["when"]
        sdf_cond = edge_data["sdf_cond"]

        fanout = [
            (v, k, edge_data)
            for _, v, k, edge_data in Graph.out_edges(node, data=True, keys=True)
        ]
        fanout_net = fanout[0][2]["name"]
        # prevent changes to self.Parasitics
        parasitic = copy.deepcopy(self.Parasitics[fanout_net])

        # Normal logical timing arc
        if unate != "rising_edge":
            type = "combinational"
            for i, out_rf in enumerate(["r", "f"]):
                if unate == "positive_unate":
                    j = i
                else:
                    j = 1 - i
                # update parasitic
                parasitic = self._Update_parasitic_(parasitic, out_rf, ECO, local_cells)
                # in_rf is relay on unate and out_rf
                Delay_LUT, Trans_LUT = (
                    self.NLDM[self.corner][cellname].delay[
                        (arc, out_rf, unate, type, when, sdf_cond)
                    ],
                    self.NLDM[self.corner][cellname].trans[
                        (arc, out_rf, unate, type, when, sdf_cond)
                    ],
                )
                order = 2
                while True:
                    try:
                        (
                            ceff_max,
                            delay_max,
                            out_trans_max,
                            fanout_trans_max,
                            fanout_delay_max,
                        ) = self.Calcu_Pin_Ceff_Delay_Trans(
                            parasitic,
                            pred_trans[j],
                            Delay_LUT,
                            Trans_LUT,
                            order,
                            cal_net=arnoldi,
                        )
                        (
                            ceff_min,
                            delay_min,
                            out_trans_min,
                            fanout_trans_min,
                            fanout_delay_min,
                        ) = self.Calcu_Pin_Ceff_Delay_Trans(
                            parasitic,
                            pred_trans[j + 2],
                            Delay_LUT,
                            Trans_LUT,
                            order,
                            cal_net=arnoldi,
                        )
                        break
                    except np.linalg.LinAlgError as e:
                        order -= 1
                        # print(f"LinAlgError occurred with order={order}, trying lower order...")

                arrival_max = delay_max + pred_arrival[j]
                arrival_min = delay_min + pred_arrival[j + 2]

                Graph[predecessor][node][key]["delay"][i] = delay_max
                Graph[predecessor][node][key]["delay"][i + 2] = delay_min
                Graph[predecessor][node][key]["trans"][i] = out_trans_max
                Graph[predecessor][node][key]["trans"][i + 2] = out_trans_min
                """
                Concurrently update arrival time, trans of node and the arrival and transition time of successors
                """
                if arnoldi:
                    if arrival_max > node_arrival[i]:
                        node_arrival[i] = arrival_max
                        for v, k, _ in fanout:
                            Graph.nodes[v]["arrival"][i] = (
                                fanout_delay_max[v] + arrival_max
                            )
                            Graph[node][v][k]["delay"][i] = fanout_delay_max[v]

                    if arrival_min < node_arrival[i + 2]:
                        node_arrival[i + 2] = arrival_min
                        for v, k, _ in fanout:
                            Graph.nodes[v]["arrival"][i + 2] = (
                                fanout_delay_min[v] + arrival_min
                            )
                            Graph[node][v][k]["delay"][i + 2] = fanout_delay_min[v]

                    if out_trans_max > node_trans[i]:
                        node_trans[i] = out_trans_max
                        for v, _, _ in fanout:
                            Graph.nodes[v]["trans"][i] = fanout_trans_max[v]

                    if out_trans_min < node_trans[i + 2]:
                        node_trans[i + 2] = out_trans_min
                        for v, _, _ in fanout:
                            Graph.nodes[v]["trans"][i + 2] = fanout_trans_min[v]
                else:
                    node_arrival[i] = max(node_arrival[i], arrival_max)
                    node_arrival[i + 2] = min(node_arrival[i + 2], arrival_min)
                    node_trans[i] = max(node_trans[i], out_trans_max)
                    node_trans[i + 2] = min(node_trans[i + 2], out_trans_min)

        # Clock rising edge timing arc"""
        else:
            type = unate
            unate = "non_unate"
            Delay_Rise_LUT, Trans_Rise_LUT = (
                self.NLDM[self.corner][cellname].delay[
                    (arc, "r", unate, type, when, sdf_cond)
                ],
                self.NLDM[self.corner][cellname].trans[
                    (arc, "r", unate, type, when, sdf_cond)
                ],
            )
            Delay_Fall_LUT, Trans_Fall_LUT = (
                self.NLDM[self.corner][cellname].delay[
                    (arc, "f", unate, type, when, sdf_cond)
                ],
                self.NLDM[self.corner][cellname].trans[
                    (arc, "f", unate, type, when, sdf_cond)
                ],
            )
            order = 2
            while True:
                try:
                    parasitic = self._Update_parasitic_(
                        parasitic, "r", ECO, local_cells
                    )
                    (
                        ceff_max_rise,
                        delay_max_rise,
                        out_trans_max_rise,
                        fanout_trans_max_rise,
                        fanout_delay_max_rise,
                    ) = self.Calcu_Pin_Ceff_Delay_Trans(
                        parasitic,
                        pred_trans[0],
                        Delay_Rise_LUT,
                        Trans_Rise_LUT,
                        order,
                        cal_net=arnoldi,
                    )
                    (
                        ceff_min_rise,
                        delay_min_rise,
                        out_trans_min_rise,
                        fanout_trans_min_rise,
                        fanout_delay_min_rise,
                    ) = self.Calcu_Pin_Ceff_Delay_Trans(
                        parasitic,
                        pred_trans[2],
                        Delay_Rise_LUT,
                        Trans_Rise_LUT,
                        order,
                        cal_net=arnoldi,
                    )
                    parasitic = self._Update_parasitic_(
                        parasitic, "f", ECO, local_cells
                    )
                    (
                        ceff_max_fall,
                        delay_max_fall,
                        out_trans_max_fall,
                        fanout_trans_max_fall,
                        fanout_delay_max_fall,
                    ) = self.Calcu_Pin_Ceff_Delay_Trans(
                        parasitic,
                        pred_trans[0],
                        Delay_Fall_LUT,
                        Trans_Fall_LUT,
                        order,
                        cal_net=arnoldi,
                    )
                    (
                        ceff_min_fall,
                        delay_min_fall,
                        out_trans_min_fall,
                        fanout_trans_min_fall,
                        fanout_delay_min_fall,
                    ) = self.Calcu_Pin_Ceff_Delay_Trans(
                        parasitic,
                        pred_trans[2],
                        Delay_Fall_LUT,
                        Trans_Fall_LUT,
                        order,
                        cal_net=arnoldi,
                    )
                    break
                except np.linalg.LinAlgError as e:
                    order -= 1
                    # print(f"LinAlgError occurred with order={order}, trying lower order...")

            arrival_max_rise = delay_max_rise + pred_arrival[0]
            arrival_max_fall = delay_max_fall + pred_arrival[0]
            arrival_min_rise = delay_min_rise + pred_arrival[2]
            arrival_min_fall = delay_min_fall + pred_arrival[2]

            Graph[predecessor][node][key]["delay"] = [
                delay_max_rise,
                delay_max_fall,
                delay_min_rise,
                delay_min_fall,
            ]
            Graph[predecessor][node][key]["trans"] = [
                out_trans_max_rise,
                out_trans_max_fall,
                out_trans_min_rise,
                out_trans_min_fall,
            ]
            """
            Concurrently update arrival time, trans of node and the arrival and transition time of successors
            """
            if arnoldi:
                if arrival_max_rise > node_arrival[0]:
                    node_arrival[0] = arrival_max_rise
                    for v, k, _ in fanout:
                        Graph.nodes[v]["arrival"][0] = (
                            fanout_delay_max_rise[v] + arrival_max_rise
                        )
                        Graph[node][v][k]["delay"][0] = fanout_delay_max_rise[v]

                if arrival_max_fall > node_arrival[1]:
                    node_arrival[1] = arrival_max_fall
                    for v, k, _ in fanout:
                        Graph.nodes[v]["arrival"][1] = (
                            fanout_delay_max_fall[v] + arrival_max_fall
                        )
                        Graph[node][v][k]["delay"][1] = fanout_delay_max_fall[v]

                if arrival_min_rise < node_arrival[2]:
                    node_arrival[2] = arrival_min_rise
                    for v, k, _ in fanout:
                        Graph.nodes[v]["arrival"][2] = (
                            fanout_delay_min_rise[v] + arrival_min_rise
                        )
                        Graph[node][v][k]["delay"][2] = fanout_delay_min_rise[v]

                if arrival_min_fall < node_arrival[3]:
                    node_arrival[3] = arrival_min_fall
                    for v, k, _ in fanout:
                        Graph.nodes[v]["arrival"][3] = (
                            fanout_delay_min_fall[v] + arrival_min_fall
                        )
                        Graph[node][v][k]["delay"][3] = fanout_delay_min_fall[v]

                if out_trans_max_rise > node_trans[0]:
                    node_trans[0] = out_trans_max_rise
                    for v, _, _ in fanout:
                        Graph.nodes[v]["trans"][0] = fanout_trans_max_rise[v]

                if out_trans_max_fall > node_trans[1]:
                    node_trans[1] = out_trans_max_fall
                    for v, _, _ in fanout:
                        Graph.nodes[v]["trans"][1] = fanout_trans_max_fall[v]

                if out_trans_min_rise < node_trans[2]:
                    node_trans[2] = out_trans_min_rise
                    for v, _, _ in fanout:
                        Graph.nodes[v]["trans"][2] = fanout_trans_min_rise[v]

                if out_trans_min_fall < node_trans[3]:
                    node_trans[3] = out_trans_min_fall
                    for v, _, _ in fanout:
                        Graph.nodes[v]["trans"][3] = fanout_trans_min_fall[v]
            else:
                node_arrival[0] = max(node_arrival[0], arrival_max_rise)
                node_arrival[1] = max(node_arrival[1], arrival_max_fall)
                node_arrival[2] = min(node_arrival[2], arrival_min_rise)
                node_arrival[3] = min(node_arrival[3], arrival_min_fall)

                node_trans[0] = max(node_trans[0], out_trans_max_rise)
                node_trans[1] = max(node_trans[1], out_trans_max_fall)
                node_trans[2] = min(node_trans[2], out_trans_min_rise)
                node_trans[3] = min(node_trans[3], out_trans_min_fall)

    # update sink caps based on rise/fall transition direction
    def _Update_parasitic_(self, parasitic, rf, ECO, local_cells):
        for i in range(len(parasitic.sink_cell_caps)):
            # endpoints_name (name, IO)
            endpoint = parasitic.endpoints_name[i + 1][0]
            # prevent mistake PO as cell
            if not ECO:
                cellname = self.cells.get(endpoint.split("/")[0], None)
            else:
                cellname = local_cells[endpoint.split("/")[0]]
            if cellname is not None:
                parasitic.sink_cell_caps[i] = self.NLDM[self.corner][
                    cellname
                ].receive_cap[(endpoint.split("/")[1], rf)]
        return parasitic

    def Estimate_ECO_Simple(self, target_node, ori_cell, new_cell):
        # new_cell is None represent Estimate current cell
        if new_cell is not None:
            candidate_cells = self.Get_Candidate(ori_cell)
            assert (
                new_cell in candidate_cells
            ), f"The footprint of {new_cell} does not match {ori_cell}"

        Copied_G, local_cells = self._Extract_Local_Subgraph(target_node)
        if new_cell is not None:
            local_cells[target_node.split("/")[0]] = new_cell
        local_levelized_node = list(nx.topological_sort(Copied_G))

        # Update timing in copied graph
        for node in local_levelized_node:
            if all(x is None for x in Copied_G.nodes[node]["arrival"]):
                continue
            in_edges = list(Copied_G.in_edges(node, data=True, keys=True))
            if not in_edges:
                pass
            else:
                node_arrival = [
                    -float("inf"),
                    -float("inf"),
                    float("inf"),
                    float("inf"),
                ]
                node_trans = [
                    -float("inf"), 
                    -float("inf"),
                    float("inf"),
                    float("inf")
                ]
                for predecessor, _, key, edge_data in Copied_G.in_edges(
                    node, data=True, keys=True
                ):
                    if (
                        edge_data.get("is_cell")
                        and not edge_data["sense_unate"].startswith("setup")
                        and not edge_data["sense_unate"].startswith("hold")
                    ):
                        self._Calc_Cell_Timing_Simple(
                            node_arrival,
                            node_trans,
                            node,
                            predecessor,
                            key,
                            edge_data,
                            ECO=True,
                            local_graph=Copied_G,
                            local_cells=local_cells,
                        )
                    elif(
                        not edge_data.get("is_cell")
                    ):
                        # # Elmore delay
                        # pred_arrival = Copied_G.nodes[predecessor][
                        #     "arrival"
                        # ]  # Max_Rise, Max_Fall, Min_Rise, Min_Fall
                        # pred_trans = Copied_G.nodes[predecessor][
                        #     "trans"
                        # ]  # Max_Rise, Max_Fall, Min_Rise, Min_Fall
                        # # not update timing with None in predecessor
                        # if None in pred_arrival:
                        #     node_arrival = Copied_G.nodes[node]["arrival"]
                        #     node_trans = Copied_G.nodes[node]["trans"]
                        # else:
                        #     for i, (ap, tp) in enumerate(
                        #         zip(pred_arrival, pred_trans)
                        #     ):
                        #         delay, trans = self.Calc_Elmore(
                        #             edge_data.get("name"), tp, node
                        #         )
                        #         Copied_G[predecessor][node][key]["delay"][i] = delay
                        #         if i < 2:
                        #             node_arrival[i] = max(
                        #                 node_arrival[i], ap + delay
                        #             )
                        #             node_trans[i] = max(node_trans[i], trans)
                        #         else:
                        #             node_arrival[i] = min(
                        #                 node_arrival[i], ap + delay
                        #             )
                        #             node_trans[i] = min(node_trans[i], trans)

                        # if the timing arc is not a cell arc, its a net arc
                        # use the timing data from PrimeaTime to calculate arrival time and transition
                        # adjust if None in predecessor to avoid clk port
                        edge_arrival = [
                            (a if a is not None else 0) + (b if b is not None else 0)
                            for a, b in zip(
                                Copied_G.nodes[predecessor]["arrival"],
                                Copied_G[predecessor][node][key]["delay"],
                            )
                        ]
                        edge_trans = [
                            a
                            - (b if b is not None else 0)
                            + (c if c is not None else 0)
                            for a, b, c in zip(
                                Copied_G.nodes[node]["trans_pt"],
                                Copied_G.nodes[predecessor]["trans_pt"],
                                Copied_G.nodes[predecessor]["trans"],
                            )
                        ]
                        # update node timing
                        for idx in [0, 1]:  # Max_Rise, Max_Fall
                            node_arrival[idx] = max(
                                node_arrival[idx], edge_arrival[idx]
                            )
                            node_trans[idx] = max(node_trans[idx], edge_trans[idx])
                        for idx in [2, 3]:  # Min_Rise, Min_Fall
                            node_arrival[idx] = min(
                                node_arrival[idx], edge_arrival[idx]
                            )
                            node_trans[idx] = min(node_trans[idx], edge_trans[idx])
                        
                Copied_G.nodes[node]["arrival"] = node_arrival
                Copied_G.nodes[node]["trans"] = node_trans
        return Copied_G
    
    def Estimate_ECO(self, target_node, ori_cell, new_cell, arnoldi=True):
        # new_cell is None represent Estimate current cell
        if new_cell is not None:
            candidate_cells = self.Get_Candidate(ori_cell)
            assert (
                new_cell in candidate_cells
            ), f"The footprint of {new_cell} does not match {ori_cell}"

        Copied_G, local_cells = self._Extract_Local_Subgraph(target_node)
        if new_cell is not None:
            local_cells[target_node.split("/")[0]] = new_cell
        local_levelized_node = list(nx.topological_sort(Copied_G))

        # Update timing in copied graph
        for node in local_levelized_node:
            if all(x is None for x in Copied_G.nodes[node]["arrival"]):
                continue
            in_edges = list(Copied_G.in_edges(node, data=True, keys=True))
            if not in_edges:
                pass
            else:
                node_arrival = [
                    -float("inf"),
                    -float("inf"),
                    float("inf"),
                    float("inf"),
                ]
                node_trans = [
                    -float("inf"), 
                    -float("inf"), 
                    float("inf"), 
                    float("inf")
                ]
                for predecessor, _, key, edge_data in Copied_G.in_edges(
                    node, data=True, keys=True
                ):
                    if (
                        edge_data.get("is_cell")
                        and not edge_data["sense_unate"].startswith("setup")
                        and not edge_data["sense_unate"].startswith("hold")
                    ):
                        self._Calc_Cell_Timing(
                            node_arrival,
                            node_trans,
                            node,
                            predecessor,
                            key,
                            edge_data,
                            arnoldi,
                            ECO=True,
                            local_graph=Copied_G,
                            local_cells=local_cells,
                        )
                    elif(
                        not edge_data.get("is_cell")
                    ):
                        # Elmore delay
                        if not arnoldi:
                            pred_arrival = Copied_G.nodes[predecessor][
                                "arrival"
                            ]  # Max_Rise, Max_Fall, Min_Rise, Min_Fall
                            pred_trans = Copied_G.nodes[predecessor][
                                "trans"
                            ]  # Max_Rise, Max_Fall, Min_Rise, Min_Fall
                            # not update timing with None in predecessor
                            if None in pred_arrival:
                                node_arrival = Copied_G.nodes[node]["arrival"]
                                node_trans = Copied_G.nodes[node]["trans"]
                            else:
                                for i, (ap, tp) in enumerate(
                                    zip(pred_arrival, pred_trans)
                                ):
                                    delay, trans = self.Calc_Elmore(
                                        edge_data.get("name"), tp, node
                                    )
                                    Copied_G[predecessor][node][key]["delay"][i] = delay
                                    if i < 2:
                                        node_arrival[i] = max(
                                            node_arrival[i], ap + delay
                                        )
                                        node_trans[i] = max(node_trans[i], trans)
                                    else:
                                        node_arrival[i] = min(
                                            node_arrival[i], ap + delay
                                        )
                                        node_trans[i] = min(node_trans[i], trans)

                if edge_data.get("is_cell") or not arnoldi:
                    Copied_G.nodes[node]["arrival"] = node_arrival
                    Copied_G.nodes[node]["trans"] = node_trans
        return Copied_G

    def _Extract_Local_Subgraph(self, node):
        Copied_G = nx.MultiDiGraph()
        local_cells = {}
        Copied_G.add_node(node, **copy.deepcopy(self.G.nodes[node]))
        # one loop backward neighbor hood, must be logical cell input
        for pred1, _, key1, edge_data1 in self.G.in_edges(node, data=True, keys=True):
            if pred1 not in Copied_G:
                Copied_G.add_node(pred1, **copy.deepcopy(self.G.nodes[pred1]))
            if not Copied_G.has_edge(pred1, node, key1):
                Copied_G.add_edge(pred1, node, key=key1, **copy.deepcopy(edge_data1))
            if pred1.split("/")[0] not in local_cells.keys():
                local_cells[pred1.split("/")[0]] = self.cells.get(
                    pred1.split("/")[0], None
                )

            # two loop backward neighbor hood, can be cell output or PI
            for pred2, _, key2, edge_data2 in self.G.in_edges(
                pred1, data=True, keys=True
            ):
                if pred2 not in Copied_G:
                    Copied_G.add_node(pred2, **copy.deepcopy(self.G.nodes[pred2]))
                if not Copied_G.has_edge(pred2, pred1, key2):
                    Copied_G.add_edge(
                        pred2, pred1, key=key2, **copy.deepcopy(edge_data2)
                    )
                if pred2.split("/")[0] not in local_cells.keys():
                    local_cells[pred2.split("/")[0]] = self.cells.get(
                        pred2.split("/")[0], None
                    )

                # three loop backward neighbor hood, can be flip flop input, or None for the first cell connect to PI
                for pred3, _, key3, edge_data3 in self.G.in_edges(
                    pred2, data=True, keys=True
                ):
                    if pred3 not in Copied_G:
                        Copied_G.add_node(pred3, **copy.deepcopy(self.G.nodes[pred3]))
                    if not Copied_G.has_edge(pred3, pred2, key3):
                        Copied_G.add_edge(
                            pred3, pred2, key=key3, **copy.deepcopy(edge_data3)
                        )
                    if pred3.split("/")[0] not in local_cells.keys():
                        local_cells[pred3.split("/")[0]] = self.cells.get(
                            pred3.split("/")[0]
                        )

                for _, succ, key4, edge_data4 in self.G.out_edges(
                    pred2, data=True, keys=True
                ):
                    if succ not in Copied_G:
                        Copied_G.add_node(succ, **copy.deepcopy(self.G.nodes[succ]))
                    if not Copied_G.has_edge(pred2, succ, key4):
                        Copied_G.add_edge(
                            pred2, succ, key=key4, **copy.deepcopy(edge_data4)
                        )
                    if succ.split("/")[0] not in local_cells.keys():
                        local_cells[succ.split("/")[0]] = self.cells.get(
                            succ.split("/")[0], None
                        )

        # one loop forward neighbor hood, must be cell input or PO
        for _, succ1, key1, edge_data1 in self.G.out_edges(node, data=True, keys=True):
            if succ1 not in Copied_G:
                Copied_G.add_node(succ1, **copy.deepcopy(self.G.nodes[succ1]))
            if not Copied_G.has_edge(node, succ1, key1):
                Copied_G.add_edge(node, succ1, key=key1, **copy.deepcopy(edge_data1))
            if succ1.split("/")[0] not in local_cells.keys():
                local_cells[succ1.split("/")[0]] = self.cells.get(
                    succ1.split("/")[0], None
                )

            # two loop forward neighbor hood, must be cell output
            for _, succ2, key2, edge_data2 in self.G.out_edges(
                succ1, data=True, keys=True
            ):
                if succ2 not in Copied_G:
                    Copied_G.add_node(succ2, **copy.deepcopy(self.G.nodes[succ2]))
                if not Copied_G.has_edge(succ1, succ2, key2):
                    Copied_G.add_edge(
                        succ1, succ2, key=key2, **copy.deepcopy(edge_data2)
                    )
                if succ2.split("/")[0] not in local_cells.keys():
                    local_cells[succ2.split("/")[0]] = self.cells.get(
                        succ2.split("/")[0], None
                    )

                # three loop forward neighbor hood, must be cell input
                for _, succ3, key3, edge_data3 in self.G.out_edges(
                    succ2, data=True, keys=True
                ):
                    if succ3 not in Copied_G:
                        Copied_G.add_node(succ3, **copy.deepcopy(self.G.nodes[succ3]))
                    if not Copied_G.has_edge(succ2, succ3, key3):
                        Copied_G.add_edge(
                            succ2, succ3, key=key3, **copy.deepcopy(edge_data3)
                        )
                    if succ3.split("/")[0] not in local_cells.keys():
                        local_cells[succ3.split("/")[0]] = self.cells.get(
                            succ3.split("/")[0], None
                        )

        return Copied_G, local_cells

    def Implement_ECO(self, Copied_G):
        # Merge timing info to self.G
        for n, ndata in Copied_G.nodes(data=True):
            self.G.nodes[n].update(copy.deepcopy(ndata))

        for u, v, key, edata in Copied_G.edges(keys=True, data=True):
            self.G[u][v][key].update(copy.deepcopy(edata))

    def Backward_Propagation(self):
        print("Propagating slack...")
        for node in reversed(self.levelized_nodes):  # propogate from PI to PO
            out_edges = list(self.G.out_edges(node, data=True, keys=True))
            if not out_edges:
                # if this node is endpoint
                if self.G.nodes[node]["is_port"] == True:
                    # if this node is a port, use its output delay
                    # slack_pt = self.G.nodes[node]["slack_pt"]
                    self.G.nodes[node]["slack"][0] = (
                        self.clk_period - self.G.nodes[node]["arrival"][0]
                    )
                    self.G.nodes[node]["slack"][1] = (
                        self.clk_period - self.G.nodes[node]["arrival"][1]
                    )
                    self.G.nodes[node]["slack"][2] = self.G.nodes[node]["arrival"][2]
                    self.G.nodes[node]["slack"][3] = self.G.nodes[node]["arrival"][3]
                    pass
                else:
                    # if this node is a data pin
                    for predecessor, _, k, edge_data in self.G.in_edges(
                        node, data=True, keys=True
                    ):
                        if (
                            self.G.nodes[predecessor]["is_clock_pin"] == True
                        ):  # only process D to CLK pin
                            # Max_Rise, Max_Fall, Min_Rise, Min_Fall
                            pred_arrival = self.G.nodes[predecessor]["arrival"]
                            pred_trans = self.G.nodes[predecessor]["trans"]
                            node_trans = self.G.nodes[node]["trans"]
                            # constraint = edge_data["delay"]

                            cellname = self.cells[
                                edge_data["name"]
                            ]  # get cell name in based on the refname
                            arc = predecessor.split("/")[1] + "->" + node.split("/")[1]
                            unate = edge_data["sense_unate"]
                            when = edge_data["when"]
                            sdf_cond = edge_data["sdf_cond"]

                            if (
                                unate.split("_")[-1] == "rise"
                            ):  # DFF is triggered by rise edge
                                # unate = "*"
                                if unate.split("_")[0] == "setup":
                                    unate = "*"
                                    constraint_max_rise = self.NLDM_Constraint(
                                        self.corner,
                                        cellname,
                                        arc,
                                        "r",
                                        unate,
                                        "setup_rising",
                                        when,
                                        sdf_cond,
                                        # if timing_type is setup_rising, we should use the CLK's min_rise trans and at
                                        pred_trans[2],
                                        node_trans[0],
                                    )

                                    constraint_max_fall = self.NLDM_Constraint(
                                        self.corner,
                                        cellname,
                                        arc,
                                        "f",
                                        unate,
                                        "setup_rising",
                                        when,
                                        sdf_cond,
                                        # if timing_type is setup_rising, we should use the CLK's min_rise trans and at
                                        pred_trans[2],
                                        node_trans[1],
                                    )

                                    self.G.nodes[node]["slack"][0] = (
                                        pred_arrival[2]
                                        + self.clk_period
                                        - constraint_max_rise
                                        - self.G.nodes[node]["arrival"][0]
                                    )
                                    self.G.nodes[node]["slack"][1] = (
                                        pred_arrival[2]
                                        + self.clk_period
                                        - constraint_max_fall
                                        - self.G.nodes[node]["arrival"][1]
                                    )
                                    """slack_pt = self.G.nodes[node]["slack_pt"]
                                    if (
                                        max(
                                            abs(
                                                np.array(slack_pt)
                                                - np.array(self.G.nodes[node]["slack"])
                                            )
                                            / (np.abs(slack_pt) + epsilon)
                                            > 0.5
                                        )
                                        > 0.5
                                    ):
                                        print(
                                            f"Warning: {node} slack changed more than 50% from pt, please check the timing data"
                                        )"""
                                    pass

                                elif unate.split("_")[0] == "hold":
                                    unate = "*"
                                    constraint_min_rise = self.NLDM_Constraint(
                                        self.corner,
                                        cellname,
                                        arc,
                                        "r",
                                        unate,
                                        "hold_rising",
                                        when,
                                        sdf_cond,
                                        # if timing_type is setup_rising, we should use the CLK's min_rise trans and at
                                        pred_trans[2],
                                        node_trans[0],
                                    )

                                    constraint_min_fall = self.NLDM_Constraint(
                                        self.corner,
                                        cellname,
                                        arc,
                                        "f",
                                        unate,
                                        "hold_rising",
                                        when,
                                        sdf_cond,
                                        # if timing_type is setup_rising, we should use the CLK's min_rise trans and at
                                        pred_trans[2],
                                        node_trans[1],
                                    )

                                    self.G.nodes[node]["slack"][2] = (
                                        self.G.nodes[node]["arrival"][2]
                                        - pred_arrival[0]
                                        - constraint_min_rise
                                    )
                                    self.G.nodes[node]["slack"][3] = (
                                        self.G.nodes[node]["arrival"][3]
                                        - pred_arrival[0]
                                        - constraint_min_fall
                                    )
                                    """slack_pt = self.G.nodes[node]["slack_pt"]
                                    if (
                                        max(
                                            abs(
                                                np.array(slack_pt)
                                                - np.array(self.G.nodes[node]["slack"])
                                            )
                                            / (np.abs(slack_pt) + epsilon)
                                            > 0.5
                                        )
                                        > 0.5
                                    ):
                                        print(
                                            f"Warning: {node} slack changed more than 50% from pt, please check the timing data"
                                        )"""
                                    pass

                            else:  # DFF is triggered by fall edge
                                # TODO: cannot process this situation now
                                pass
            else:
                # if this node is on data path
                node_slack = [float("inf"), float("inf"), float("inf"), float("inf")]
                for _, successor, k, edge_data in out_edges:
                    succ_slack = self.G.nodes[successor]["slack"]
                    succ_at = self.G.nodes[successor]["arrival"]
                    edge_delay = edge_data["delay"]
                    unate = edge_data["sense_unate"]

                    for i, out_rf in enumerate(["r", "f"]):
                        if unate == "negative_unate":
                            j = 1 - i
                        else:
                            j = i

                        if (
                            succ_slack[j] != float("inf")
                            and succ_slack[j] is not None
                            and succ_at[j] is not None
                            and edge_delay[j] is not None
                            and self.G.nodes[node]["arrival"][i] is not None
                        ):
                            node_slack[i] = min(
                                node_slack[i],
                                succ_slack[j]
                                + succ_at[j]
                                - edge_delay[j]
                                - self.G.nodes[node]["arrival"][i],
                            )

                        if (
                            succ_slack[j + 2] != float("inf")
                            and succ_slack[j + 2] is not None
                            and succ_at[j + 2] is not None
                            and edge_delay[j + 2] is not None
                            and self.G.nodes[node]["arrival"][i + 2] is not None
                        ):
                            node_slack[i + 2] = min(
                                node_slack[i + 2],
                                self.G.nodes[node]["arrival"][i + 2]
                                + succ_slack[j + 2]
                                + edge_delay[j + 2]
                                - succ_at[j + 2],
                            )
                self.G.nodes[node]["slack"] = node_slack

                """
                slack_pt = self.G.nodes[node]["slack_pt"]
                if (
                    max(
                        abs(
                            np.array(slack_pt)
                            - np.array(self.G.nodes[node]["slack"])
                        )
                        / (np.abs(slack_pt) + epsilon)
                        > 0.5
                    )
                    > 0.5
                ):
                    print(
                        f"Warning: {node} slack changed more than 50% from pt, please check the timing data"
                    )"""
                pass

    def Report_Global_Timing(self):
        TNS = [0, 0, 0, 0]
        WNS = [0, 0, 0, 0]
        for node in self.G.nodes():
            slack = self.G.nodes[node].get("slack")
            if len(self.G.out_edges(node)) == 0:
                for idx in range(4):
                    if slack[idx] is not None and slack[idx] < 0:
                        TNS[idx] += slack[idx]
                        WNS[idx] = min(WNS[idx], slack[idx])
        print("Global Timing Report:")
        print("\t\tTNS\t\t\tWNS")
        print(f"setup_rise:\t{TNS[0]:.4f}\t\t\t{WNS[0]:.4f}")
        print(f"setup_fall:\t{TNS[1]:.4f}\t\t\t{WNS[1]:.4f}")
        print(f"hold_rise:\t{TNS[2]:.4f}\t\t\t{WNS[2]:.4f}")
        print(f"hold_fall:\t{TNS[3]:.4f}\t\t\t{WNS[3]:.4f}")


if __name__ == "__main__":
    TP = Timing_Propagation(design="mc_top")
    TP.clock_period = 0.5
    # TP.ECO_Iteration(max_iteration=5, Simple=False)
    # TP.Estimate_ECO('FE_OFC658_n6151/ZN', 'INVD1BWP16P90ULVT', 'INVSKPD10BWP16P90ULVT')
    # TP.Forward_Propagation_Simple(ECO=True)
    TP.Forward_Propagation(ECO=False)
    # TP.Backward_Propagation()
    # TP.Report_Global_Timing()

    from sklearn.metrics import r2_score

    arrival_all = []
    arrival_pt_all = []
    trans_all = []
    trans_pt_all = []
    delay_all = []
    delay_pt_all = []
    slack_all = []
    slack_pt_all = []

    for node in TP.levelized_nodes:
        arrival = TP.G.nodes[node].get("arrival")
        arrival_pt = TP.G.nodes[node].get("arrival_pt")
        arrival = [0.0 if v is None else v for v in arrival]
        arrival_pt = [0.0 if v is None else v for v in arrival_pt]
        arrival_all.extend(arrival)
        arrival_pt_all.extend(arrival_pt)

        trans = TP.G.nodes[node].get("trans")
        trans_pt = TP.G.nodes[node].get("trans_pt")
        trans = [0.0 if v is None else v for v in trans]
        trans_pt = [0.0 if v is None else v for v in trans_pt]
        trans_all.extend(trans)
        trans_pt_all.extend(trans_pt)

        slack = TP.G.nodes[node].get("slack")
        slack_pt = TP.G.nodes[node].get("slack_pt")

        for i in range(4):
            if slack_pt[i] is None:
                slack_pt[i] = 0.0
                slack[i] = 0.0
            else:
                slack_pt[i] = slack_pt[i]
                slack[i] = slack[i]

        # print(slack)
        # print(slack_pt)
        slack_all.extend(slack)
        slack_pt_all.extend(slack_pt)

        # print(node, arrival, arrival_pt, trans, trans_pt)

    for u, v, k, data in TP.G.edges(keys=True, data=True):
        delay_all.extend(data["delay"])
        delay_pt_all.extend(data["delay_pt"])

    arrival_all = np.array(arrival_all)
    arrival_pt_all = np.array(arrival_pt_all)
    trans_all = np.array(trans_all)
    trans_pt_all = np.array(trans_pt_all)
    delay_all = np.array(delay_all)
    delay_pt_all = np.array(delay_pt_all)
    slack_all = np.array(slack_all)
    slack_pt_all = np.array(slack_pt_all)
    r2_at = r2_score(arrival_pt_all, arrival_all)
    r2_tr = r2_score(trans_pt_all, trans_all)
    r2_dl = r2_score(delay_pt_all, delay_all)
    r2_sl = r2_score(slack_pt_all, slack_all)
    print(f"整体 R² 相似度为: arrival_time {r2_at:.8f} transition_time {r2_tr:.8f} delay {r2_dl:.8f} slack {r2_sl:.8f}")

    epsilon = 1e-100
    percentage_abs_error = np.abs(arrival_pt_all - arrival_all) / (np.abs(arrival_pt_all) + epsilon) * 100
    avg_pae_at = np.mean(percentage_abs_error)

    percentage_abs_error = np.abs(trans_pt_all - trans_all) / (np.abs(trans_pt_all) + epsilon) * 100
    avg_pae_tr = np.mean(percentage_abs_error)

    percentage_abs_error = np.abs(delay_pt_all - delay_all) / (np.abs(delay_pt_all) + epsilon) * 100
    avg_pae_dl = np.mean(percentage_abs_error)

    print(f"平均百分比绝对误差MPAE为: arrival_time {avg_pae_at:.4f}% transition_time {avg_pae_tr:.4f}% delay {avg_pae_dl:.4f}%")

    # for u, v, k, data in TP.G.edges(keys=True, data=True):
    #     delay = np.array(data['delay'], dtype=float)
    #     delay_pt = np.array(data['delay_pt'], dtype=float)

    #     denominator = np.abs(delay_pt) + 1e-8
    #     relative_error = np.abs(delay - delay_pt) / denominator

    #     if np.any(relative_error > 0.1):
    #         print(f"{u} -> {v} (key={k})")
    #         print(f"input trans: {TP.G.nodes[u]["trans"]}")
    #         print(f"input_pt: {TP.G.nodes[u]["trans_pt"]}")
    #         print(f"  delay    : {delay}")
    #         print(f"  delay_pt : {delay_pt}")
    #         print(f"  rel_error: {relative_error}")
